package me.lwn.auth.config;

import com.nimbusds.jose.jwk.JWKSet;
import com.nimbusds.jose.jwk.RSAKey;
import com.nimbusds.jose.jwk.source.JWKSource;
import com.nimbusds.jose.proc.SecurityContext;
import me.lwn.auth.security.jose.Jwks;
import me.lwn.auth.security.oauth2.server.RedisOAuth2AuthorizationService;
import me.lwn.auth.security.oauth2.server.authorization.CustomOAuth2TokenCustomizer;
import me.lwn.auth.security.oauth2.server.authorization.authentication.OAuth2PasswordAuthenticationProvider;
import me.lwn.auth.security.oauth2.server.authorization.authentication.OAuth2SmsAuthenticationProvider;
import me.lwn.auth.security.oauth2.server.authorization.oidc.authentication.DefaultOidcUserInfoMapper;
import me.lwn.auth.security.oauth2.server.authorization.web.authentication.OAuth2PasswordAuthenticationConverter;
import me.lwn.auth.security.oauth2.server.authorization.web.authentication.OAuth2SmsAuthenticationConverter;
import me.lwn.auth.security.provisioning.CustomizeJdbcUserDetailsManager;
import me.lwn.auth.security.web.authentication.LoginLockedSupport;
import me.lwn.auth.security.web.authentication.TokenAuthenticationFailureHandler;
import me.lwn.auth.security.web.authentication.TokenAuthenticationSuccessHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.security.authentication.AuthenticationManager;
import org.springframework.security.config.Customizer;
import org.springframework.security.config.annotation.web.builders.HttpSecurity;
import org.springframework.security.config.annotation.web.configuration.OAuth2AuthorizationServerConfiguration;
import org.springframework.security.config.annotation.web.configurers.oauth2.server.authorization.OAuth2AuthorizationServerConfigurer;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.crypto.password.NoOpPasswordEncoder;
import org.springframework.security.oauth2.jwt.JwtDecoder;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.server.authorization.JdbcOAuth2AuthorizationConsentService;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationConsentService;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenCustomizer;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationProvider;
import org.springframework.security.oauth2.server.authorization.client.JdbcRegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.config.ProviderSettings;
import org.springframework.security.oauth2.server.authorization.web.authentication.DelegatingAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2AuthorizationCodeAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2ClientCredentialsAuthenticationConverter;
import org.springframework.security.oauth2.server.authorization.web.authentication.OAuth2RefreshTokenAuthenticationConverter;
import org.springframework.security.web.SecurityFilterChain;
import org.springframework.security.web.util.matcher.RequestMatcher;

import javax.sql.DataSource;
import java.util.Arrays;

import static me.lwn.auth.security.utils.UserDetailsSqlConstants.USERS_BY_USERNAME_QUERY;

/**
 * @author lizhichao
 * @since 0.0.1
 */
@Configuration(proxyBeanMethods = false)
public class AuthorizationServerConfig {

    private final LoginLockedSupport loginLockedSupport;

    public AuthorizationServerConfig(LoginLockedSupport loginLockedSupport) {
        this.loginLockedSupport = loginLockedSupport;
    }

    @Bean
    @Order(Ordered.HIGHEST_PRECEDENCE)
    public SecurityFilterChain authorizationServerSecurityFilterChain(HttpSecurity http, DataSource dataSource,
                                                                      RegisteredClientRepository registeredClientRepository,
                                                                      OAuth2AuthorizationService authorizationService,
                                                                      JwtDecoder jwtDecoder) throws Exception {

        UserDetailsService userDetailsService = http.getSharedObject(UserDetailsService.class);
        if (userDetailsService == null) {
            userDetailsService = new CustomizeJdbcUserDetailsManager(dataSource);
            ((CustomizeJdbcUserDetailsManager) userDetailsService).setUsersByUsernameQuery(USERS_BY_USERNAME_QUERY);
            http.setSharedObject(UserDetailsService.class, userDetailsService);
        }
        http.setSharedObject(OAuth2TokenCustomizer.class, new CustomOAuth2TokenCustomizer(userDetailsService));
        OAuth2AuthorizationServerConfigurer<HttpSecurity> authorizationServerConfigurer = new OAuth2AuthorizationServerConfigurer<>();

        http.apply(authorizationServerConfigurer.oidc(oidcConfigurer ->
                oidcConfigurer.userInfoEndpoint(oidcUserInfoEndpointConfigurer ->
                        oidcUserInfoEndpointConfigurer.userInfoMapper(new DefaultOidcUserInfoMapper()))));

        RequestMatcher endpointsMatcher = authorizationServerConfigurer.getEndpointsMatcher();

        http
                .requestMatcher(endpointsMatcher)
                .authorizeRequests(authorizeRequests -> authorizeRequests.anyRequest().authenticated())
                .csrf(csrf -> csrf.ignoringRequestMatchers(endpointsMatcher))
                .apply(authorizationServerConfigurer);

        http.apply(authorizationServerConfigurer.tokenEndpoint((tokenEndpoint) -> {

            // 扩展支持Password认证模式
            OAuth2PasswordAuthenticationConverter oAuth2PasswordAuthenticationConverter = new OAuth2PasswordAuthenticationConverter();
            oAuth2PasswordAuthenticationConverter.setLoginLocked(loginLockedSupport);

            //扩展支持SMS认证模式
            OAuth2SmsAuthenticationConverter smsAuthenticationConverter = new OAuth2SmsAuthenticationConverter();
            smsAuthenticationConverter.setLoginLocked(loginLockedSupport);

            tokenEndpoint.accessTokenRequestConverter(
                    new DelegatingAuthenticationConverter(Arrays.asList(
                            new OAuth2AuthorizationCodeAuthenticationConverter(),
                            new OAuth2RefreshTokenAuthenticationConverter(),
                            new OAuth2ClientCredentialsAuthenticationConverter(),
                            oAuth2PasswordAuthenticationConverter,
                            smsAuthenticationConverter))
            );
            TokenAuthenticationFailureHandler failureHandler = new TokenAuthenticationFailureHandler();
            failureHandler.setLoginLocked(loginLockedSupport);
            tokenEndpoint.errorResponseHandler(failureHandler);

            TokenAuthenticationSuccessHandler successHandler = new TokenAuthenticationSuccessHandler(jwtDecoder);
            successHandler.setLoginLocked(loginLockedSupport);
            tokenEndpoint.accessTokenResponseHandler(successHandler);
        }));

        http.apply(authorizationServerConfigurer.clientAuthentication(client -> {
            OAuth2ClientAuthenticationProvider clientAuthenticationProvider =
                    new OAuth2ClientAuthenticationProvider(registeredClientRepository,
                            authorizationService);
            clientAuthenticationProvider.setPasswordEncoder(NoOpPasswordEncoder.getInstance());
            client.authenticationProvider(clientAuthenticationProvider);
        }));

        http.oauth2ResourceServer().jwt();
        SecurityFilterChain securityFilterChain = http.formLogin(Customizer.withDefaults()).build();

        addOAuth2PasswordAuthenticationProvider(http);
        addOAuth2SmsAuthenticationProvider(http);

        return securityFilterChain;
    }

    //创建Password模式token解析器
    private void addOAuth2PasswordAuthenticationProvider(HttpSecurity http) {

        AuthenticationManager authenticationManager = http.getSharedObject(AuthenticationManager.class);
        ProviderSettings providerSettings = http.getSharedObject(ProviderSettings.class);
        OAuth2AuthorizationService authorizationService = http.getSharedObject(OAuth2AuthorizationService.class);
        JwtEncoder jwtEncoder = http.getSharedObject(JwtEncoder.class);
        OAuth2PasswordAuthenticationProvider passwordAuthenticationProvider =
                new OAuth2PasswordAuthenticationProvider(authenticationManager, authorizationService, jwtEncoder);

        passwordAuthenticationProvider.setProviderSettings(providerSettings);
        http.authenticationProvider(passwordAuthenticationProvider);
    }

    //创建Sms模式token解析器
    private void addOAuth2SmsAuthenticationProvider(HttpSecurity http) {

        AuthenticationManager authenticationManager = http.getSharedObject(AuthenticationManager.class);
        ProviderSettings providerSettings = http.getSharedObject(ProviderSettings.class);
        OAuth2AuthorizationService authorizationService = http.getSharedObject(OAuth2AuthorizationService.class);
        JwtEncoder jwtEncoder = http.getSharedObject(JwtEncoder.class);
        OAuth2SmsAuthenticationProvider smsAuthenticationProvider =
                new OAuth2SmsAuthenticationProvider(authenticationManager, authorizationService, jwtEncoder);

        smsAuthenticationProvider.setProviderSettings(providerSettings);
        http.authenticationProvider(smsAuthenticationProvider);
    }

    @Bean
    public RegisteredClientRepository registeredClientRepository(JdbcTemplate jdbcTemplate) {
        return new JdbcRegisteredClientRepository(jdbcTemplate);
    }

    @Bean
    public OAuth2AuthorizationService authorizationService(RedisConnectionFactory connectionFactory, RegisteredClientRepository repository) {
        return new RedisOAuth2AuthorizationService(connectionFactory, repository);
    }

    @Bean
    public OAuth2AuthorizationConsentService authorizationConsentService(JdbcTemplate jdbcTemplate, RegisteredClientRepository registeredClientRepository) {
        return new JdbcOAuth2AuthorizationConsentService(jdbcTemplate, registeredClientRepository);
    }

    @Bean
    public JWKSource<SecurityContext> jwkSource() {
        RSAKey rsaKey = Jwks.generateRsa();
        JWKSet jwkSet = new JWKSet(rsaKey);
        return (jwkSelector, context) -> jwkSelector.select(jwkSet);
    }

    @Bean
    JwtDecoder jwtDecoder(JWKSource<SecurityContext> jwkSource) {
        return OAuth2AuthorizationServerConfiguration.jwtDecoder(jwkSource);
    }

    @Bean
    public ProviderSettings providerSettings() {
        return ProviderSettings.builder().issuer("http://127.0.0.1:9000").build();
    }
}
